import numpy as np

# 计算 E 步中的 μ_j^{(i+1)}
def EM_mu_j(pi, p, q):
    list_mu = []
    list_mu.append(pi * (1 - p) / ( pi * (1 - p) + (1 - pi) * (1 - q)))
    list_mu.append(pi * p / ( pi * p + (1 - pi) * q))
    return list_mu
# 计算 E 步中所有的 μ^{(i+1)}
def EM_mu(y,pi, p, q):
    mu = []
    mu_j = EM_mu_j(pi, p, q)
    for i in range(len(y)):
        mu.append(mu_j[y[i]])
    return mu
# 计算 M 步中的 π
def EM_pi(mu):
    return np.mean(mu)
# 计算 M 步中的 p
def EM_p(y, mu):
    return np.sum(mu * y) / np.sum(mu)
# 计算 M 步中的 q
def EM_q(y, mu):
    return np.sum((1 - mu) * y) / np.sum(1 - mu)

# EM 算法迭代停止条件和最大次数
error = 1e-6
max_iterations = 10
# 观测数据与初始参数值
y = [1, 1, 0, 1, 0, 0, 1, 0, 1, 1]
pi = 0.46
p = 0.55
q = 0.67

for iteration in range(max_iterations):
    print(f"Iteration {iteration + 1}:")
    # E 步：计算 μ_j^{(i+1)}
    mu = EM_mu(y, pi, p, q)
    print(f"mu: {mu}")
    # M 步：更新参数 π, p, q
    pi_new = EM_pi(np.array(mu))
    p_new = EM_p(np.array(y), np.array(mu))
    q_new = EM_q(np.array(y), np.array(mu))
    # 判断停止条件
    if np.abs(pi_new - pi) < error and np.abs(p_new - p) < error and np.abs(q_new - q) < error:
        print("Convergence reached.")
        break
    # 打印当前参数值
    print(f"Current parameters:")
    print(f"pi: {pi_new}, p: {p_new}, q: {q_new}\n")
    # 更新参数值
    pi, p, q = pi_new, p_new, q_new

print("Final parameters:")
print(f"pi: {pi}, p: {p}, q: {q}")